#include <elf.h>
#include <link.h>
#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include <ctype.h>
#include <unistd.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/syscall.h>
#include <sys/user.h>
#include <stdint.h>
#include <inttypes.h>
#include <errno.h>
#include "elf_hook.h"


#if defined(__LP64__)
#define XH_ELF_R_SYM(info)  ELF64_R_SYM(info)
#define XH_ELF_R_TYPE(info) ELF64_R_TYPE(info)
#else
#define XH_ELF_R_SYM(info)  ELF32_R_SYM(info)
#define XH_ELF_R_TYPE(info) ELF32_R_TYPE(info)
#endif

#define PAGE_START(addr) ((addr) & PAGE_MASK)
#define PAGE_END(addr)   (PAGE_START(addr + sizeof(uintptr_t) - 1) + PAGE_SIZE)

static ElfW(Phdr)*
elf_get_first_segment_by_type(elf_t *self, ElfW(Word)type)
{
	ElfW(Phdr)* phdr;
	for(phdr = self->phdr; phdr < self->phdr + self->ehdr->e_phnum; phdr++){
		if(phdr->p_type == type){
			return phdr;
		}
	}
	return NULL;
}

int elf_check_file_header(uintptr_t base_addr)
{
	ElfW(Ehdr) *ehdr = (ElfW(Ehdr)*)base_addr;

	if(0 != memcmp(ehdr->e_ident, "\177ELF", 4)) return -1;
	if(ELFCLASS32 != ehdr->e_ident[EI_CLASS]) return -2;

	// 简单校验
	return 0;
}

int elf_init(elf_t *self, uintptr_t base_addr, const char* pathname)
{
	memset(self, 0, sizeof(elf_t));
	self->pathname = pathname;
	self->base_addr = (ElfW(Addr))base_addr;
	self->ehdr = (ElfW(Ehdr)*)base_addr;
	self->phdr = (ElfW(Phdr)*)(base_addr + self->ehdr->e_phoff);

	// find the first load-segment with offset 0
	ElfW(Phdr)* phdr0 = elf_get_first_segment_by_type_offset(self, PT_LOAD, 0);
	if(NULL == phdr0){
		abort();
	}

	if(self->base_addr < phdr0->p_vaddr) return -1;
	self->bias_addr = self->base_addr - phdr0->p_vaddr;

	ElfW(Phdr)* dhdr = elf_get_first_segment_by_type(self, PT_DYNAMIC);
	if(NULL == dhdr){
		abort();
	}

	self->dyn = (ElfW(Dyn)*)(self->bias_addr + dhdr->p_vaddr);
	self->dyn_sz = dhdr->p_memsz;

	ElfW(Dyn) *dyn = self->dyn;
	ElfW(Dyn) *dyn_end = self->dyn + (self->dyn_sz / sizeof(ElfW(Dyn)));
	uint32_t *raw;
	for(;dyn < dyn_end; dyn++)
	{
		switch(dyn->d_tag)
		{
		case DT_NULL:
			dyn = dyn_end;
			break;
		case DT_STRTAB:
		{
			self->strtab = (const char*)(self->bias_addr + dyn->d_un.d_ptr);
			if((ElfW(Addr))(self->strtab) < self->base_addr){
				abort();
			}
			break;
		}
		case DT_SYMTAB:
		{
			self->symtab = (ElfW(Sym)*)(self->bias_addr + dyn->d_un.d_ptr);
			if((ElfW(Addr))(self->symtab) < self->base_addr)
				abort();
			break;
		}
		case DT_PLTREL:
			break;
		case DT_JMPREL:
		{
			self->relplt = dyn->d_un.d_ptr;
//			self->relplt = (ElfW(Addr))(self->bias_addr + dyn->d_un.d_ptr);
			break;
		}
		case DT_PLTRELSZ:
			self->replt_sz = dyn->d_un.d_val;
			break;
		default:
			break;
		}
	}
	return 0;
}

ElfW(Phdr) *
elf_get_first_segment_by_type_offset(elf_t* self, ElfW(Word)type, ElfW(Off) offset)
{
	ElfW(Phdr) * phdr;
	for(phdr = self->phdr; phdr < self->phdr + self->ehdr->e_phnum; phdr++){
		if(phdr->p_type == type && phdr->p_offset == offset){
			return phdr;
		}
	}
	return NULL;
}

void elf_print(elf_t* self)
{
	printf("elf file: base_addr%p  bias_addr: %p\n", self->base_addr, self->bias_addr);
	printf("    .dynamic: %p   size: %d\n", self->dyn, self->dyn_sz);
	printf("    .rel.plt: %p\n", self->relplt);
}


elf_t* read_from_maps()
{
	char line[512];
	FILE *fp;
	uintptr_t base_addr;
	char perm[5];
	unsigned long offset;
	int pathname_pos;
	char *pathname;
	size_t pathname_len;
	if (NULL == (fp=fopen("/proc/self/maps", "r"))){
		abort();
	}

	while(fgets(line, sizeof(line), fp))
	{
		if(sscanf(line, "%x-%*lx %4s %lx %*x:%*x %*d%n",
				  &base_addr, perm, &offset, &pathname_pos) != 3)
			continue;
		if(perm[0] != 'r') continue;
		if(perm[3] != 'p') continue;
		if(0 != offset) continue;

		while(isspace(line[pathname_pos]) && pathname_pos < (int)(sizeof(line) - 1))
			pathname_pos += 1;
		if(pathname_pos >= (int)(sizeof(line) - 1)) continue;
		pathname = line + pathname_pos;
		pathname_len = strlen(pathname);
		if(0 == pathname_len) continue;
		if(pathname[pathname_len - 1] == '\n'){
			pathname[pathname_len - 1] = '\0';
			pathname_len -= 1;
		}
		if(0 == pathname_len) continue;
		if('[' == pathname[0]) continue;

		if(!elf_check_file_header(base_addr)
		   && strcmp("/home/ldk/sources/c/clearn/plthook/build/main", pathname) == 0){
			printf("find elf: %s\n", pathname);
			elf_t* elf = (elf_t*)malloc(sizeof(elf_t));
			if(elf_init(elf, base_addr, pathname) == 0){
				elf_print(elf);
				return elf;
			} else {
				return NULL;
			}
		}
	}

	return NULL;
}

static void elf_replace_func(elf_t* self,
							 ElfW(Addr) addr, void* new_func, void**old_func){
	void *old_addr = *(void**)addr;
	*old_func = old_addr;
	*(void** )addr = new_func;
	__builtin___clear_cache((void*)PAGE_START(addr), (void*)PAGE_END(addr));
}

int elf_hook(elf_t *self, void* new_func, void** old_func)
{
	uint32_t symidx;
	void *rel_common;

	uint8_t* start = (uint8_t *)self->relplt;
	uint8_t* end = start + self->replt_sz;

	ElfW(Rel) *rel;
	ElfW(Addr) r_offset;
	size_t r_info;
	size_t r_sym;
	size_t r_type;
	ElfW(Addr) addr;
	int r;
	while (start < end){
		rel = (ElfW(Rel)*)start;
		r_info = rel->r_info;
		r_offset = rel->r_offset;
		r_sym = XH_ELF_R_SYM(r_info);
		r_type = XH_ELF_R_TYPE(r_info);
		if(r_type != R_386_JMP_SLOT){
			continue;
		}
		if(r_offset == 0x4024){
			printf("find sym:\n");
			addr = self->bias_addr + r_offset;
			if(addr < self->base_addr){
				abort();
			}
			elf_replace_func(self, addr, new_func, old_func);
			break;
		}
		start += sizeof(ElfW(Rel));
	}
}
